from keras.preprocessing.image import ImageDataGenerator
import numpy as np


def data_augmentation(X, Y):
    # this will do preprocessing and realtime data augmentation
    datagen = ImageDataGenerator(
        featurewise_center=False,  # set input mean to 0 over the dataset
        samplewise_center=False,  # set each sample mean to 0
        featurewise_std_normalization=False,  # divide inputs by std of the dataset
        samplewise_std_normalization=False,  # divide each input by its std
        zca_whitening=False,  # apply ZCA whitening
        rotation_range=20,  # randomly rotate images in the range (degrees, 0 to 180)
        width_shift_range=0.1,  # randomly shift images horizontally (fraction of total width)
        height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
        horizontal_flip=True,  # randomly flip images
        vertical_flip=False)  # randomly flip images

    # compute quantities required for featurewise normalization
    # (std, mean, and principal components if ZCA whitening is applied)
    datagen.fit(X)
    return datagen


def normlize_per_channel(img, img_size=None, crop_size=None, color_mode="rgb", out=None):
    img[:, :, 0] -= 123.68
    img[:, :, 1] -= 116.779
    img[:, :, 2] -= 103.939
    img = img.transpose((2, 0, 1))
    img_list = []

    if crop_size:
        img = img[:, (img_size[0] - crop_size[0]) // 2:(img_size[0] + crop_size[0]) // 2
        ,(img_size[1] - crop_size[1]) // 2:(img_size[1] + crop_size[1]) // 2]

    img_list.append(img)

    img_batch = np.hstack(img_list)
    if not out is None:
        out.append(img_batch)
    else:
        return img_batch


def preprocessing(X, config):
    return X
